import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.recommendation.ALS
import org.apache.spark.sql.catalyst.dsl.expressions.StringToAttributeConversionHelper
import org.apache.spark.sql.functions.{col, concat_ws}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

object MenuMain {

  // 添加隐式转换
  implicit class DataFrameHelpers(df: DataFrame) {
    def showWithIndex(): Unit = {
      df.rdd.zipWithIndex().map{ case (row, index) => s"${index + 1}: $row"}.collect().foreach(println)
    }
  }

  def main(args: Array[String]): Unit = {
    // 创建SparkSession
    val spark = SparkSession.builder()
      .appName("MenuRecommendation")
      .master("local")
      .getOrCreate()

    // 读取数据
    val data = spark.read.option("header", "true").option("inferSchema", "true")
      .csv("menu.csv")

    // 将数据拆分为训练集和测试集
    val Array(training, test) = data.filter(x => x!=null).randomSplit(Array(0.8, 0.2))

    // 创建ALS模型
    val als = new ALS()
      .setMaxIter(5)
      .setRegParam(0.01)
      .setUserCol("userId")
      .setItemCol("dishId")
      .setRatingCol("rating")

    // 使用训练数据拟合ALS模型
    val model = als.fit(training)

    // 对测试数据进行预测
    val predictions = model.transform(test).na.drop()

    // 使用RegressionEvaluator来评估模型
    val evaluator = new RegressionEvaluator()
      .setMetricName("rmse")
      .setLabelCol("rating")
      .setPredictionCol("prediction")

    val rmse = evaluator.evaluate(predictions)
    println(s"Root-mean-square error = $rmse")

    // 为特定用户推荐菜品
    val userId = 1 // 你想要为哪个用户生成推荐
    val userRecs = model.recommendForAllUsers(3)
    // 将整数列转换为字符串格式
    val rdd = userRecs.select("userId", "recommendations").rdd.map {
      case Row(userId: Int, recommendations: Seq[Row]) =>
        s"$userId,${recommendations.map(row => s"${row.getInt(0)},${row.getFloat(1)}").mkString(";")}"
    }

    rdd.foreach(x => println(x))
    // 保存 RDD 到文本文件
//    rdd.repartition(1).saveAsTextFile("file:///D:\\\\java-project\\\\spark-menu\\\\user1")

    //      .filter(df => df.getInt(0) == userId) // 过滤特定用户的推荐结果
//    println("menu:")
//    userRecs.show()
//    userRecs.write
//      .format("text")
////      .option("header", "true")
//      .mode("overwrite") // 如果文件已存在，覆盖它
//      .save("file:///D:\\java-project\\spark-menu\\user") // 使用 file:// 协议指定本地文件路径


    // 为特定菜品推荐用户
    val dishId = 101 // 你想要为哪道菜品生成推荐
    val dishRecs = model.recommendForAllItems(3)
//      .filter(df => df.getInt(0) == dishId) // 过滤特定菜品的推荐结果
//    dishRecs.show()
    dishRecs.foreach(println(_))


//
//    dishRecs.repartition(1)
//      .write
//      .option("header", "true").save("userRes")

    // 停止SparkSession
    spark.stop()
  }
}
